#! /usr/bin/env python
# coding=utf-8

import os
from Bio import SeqIO
from Bio.Seq import Seq
from barcode_checking import *

import sys
import argparse
from argparse import RawTextHelpFormatter



parser = argparse.ArgumentParser(
    description='''
    提取 barcode_fix_seq.py 修正完成的结果
    用法:
    barcoding_extract_result.py -i ngsp_filter_out_all -o result.txt -t 5
    由大天才于2022年3月24日创建于浙江农业大学''',formatter_class=RawTextHelpFormatter)



parser.add_argument('-i',
                help='必须给定,NGSpeciesID标本的路径')

parser.add_argument('-t',
                help='翻译table，默认为无脊椎动物线粒体 table=5')

parser.add_argument('-o',
                help='必须给定，输出文件')



args = parser.parse_args()

if not args.i or not args.o:

    parser.print_help()

    sys.exit()


inpath = args.i 

outfile = args.o

if not args.t:
	table = 5
else:
	table = int(args.t)


outfile = open(outfile,'w')
outline ='\t'.join([
	'specimen_name',
	'cluster_name',
	'cluster_id',
	'reads_coverage',
	'raw_seq_length',
	'raw_seq_translatable',
	'unmatched',
	'fixed_seq_length',
	'fixed_seq_translatable',
	'mismatch_5p',
	'mismatch_3p',
	'insertion_error_number',
	'deletion_error_number',
	'insertion',
	'deletion',
	'insertion codon',
	'deletion codon',
	'raw_seq',
	'fixed_seq',
	'fixed_translation',
	])+'\n'
outfile.write(outline)



for i in os.listdir(inpath):
	p = inpath+'/'+i
	for j in os.listdir(p):
		if j.split('.')[-1] =='fasta':

			cluster_name = '.'.join(j.split('.')[:-1])
			cluster_id = i+'|'+cluster_name
			
			
			fasta_file_path = p+'/'+j
			fix_file_path = p+'/'+j+'.fixed'
			
			# 读取reads_nub
			raw_seq = next(SeqIO.parse(fasta_file_path,'fasta'))
			raw_seq_seq = str(raw_seq.seq)
			raw_seq_name = raw_seq.name
			raw_seq_coverage = int(raw_seq.name.split('_')[-1])
			raw_trans_pep = check_consensus_by_translation(raw_seq.seq,table=table)
			raw_trans = (raw_trans_pep!= False and raw_trans_pep!='')
			#print(raw_trans)

			#print(raw_seq_coverage)
			# 读取修正信息

			fix_lista = next(open(fix_file_path)).split('\t')
			fix_seq_seq = fix_lista[0]


			# 翻译修正
			fix_seq = Seq(fix_seq_seq)
			fix_trans_pep = check_consensus_by_translation(fix_seq , table=table)
			if fix_trans_pep ==False:
				fix_trans_pep=''
			fix_trans = (fix_trans_pep!= False and fix_trans_pep!='')

			#print(fix_trans)
			state = fix_lista[1]
			state_lista = state.split(';')
			
			unmatched = False
			if state.find('unmatch') !=-1:
				unmatched = True


			mismatch_3p = False
			mismatch_5p = False
			if state.find('side mismatch 3p') !=-1:
				mismatch_3p = True
			if state.find('side mismatch 5p') !=-1:
				mismatch_5p = True
			

			insertion_error = []
			deletion_error = []
			insertion_codon = []
			deletion_codon = []

			for z in state_lista:
				if z.find('insertion error')!=-1:
					insertion_error.append(z)
				elif z.find('deletion error') !=-1:
					deletion_error.append(z)
				elif z.find('deletion codon') !=-1:
					deletion_codon.append(z)
				elif z.find('insertion codon') !=-1:
					insertion_codon.append(z)

			# 输出的排序
			outline ='\t'.join([
				i,
				cluster_name,
				cluster_id,
				str(raw_seq_coverage),
				str(len(raw_seq_seq)),
				str(raw_trans),
				str(unmatched),
				str(len(fix_seq_seq)),
				str(fix_trans),
				str(mismatch_5p),
				str(mismatch_3p),

				str(len(insertion_error)),
				str(len(deletion_error)),
				
				';'.join(insertion_error),
				';'.join(deletion_error),

				';'.join(deletion_codon),
				';'.join(insertion_codon),

				str(raw_seq_seq),
				str(fix_seq_seq),
				str(fix_trans_pep),
				])+'\n'
			outfile.write(outline)


outfile.close()


